import bcolz
from pathlib import Path
import os

from pandas import to_datetime
import pandas as pd
import shutil


path = "d:/jqdata/quote"


def open_minute_file():
    global path
    paths = Path(path) / "stock_index_daily"
    write_path = Path(path) / "stock_index_dailys"
    dirs = os.listdir(paths)
    for dir in dirs:
        if dir <="2006":
            continue
        print(dir)
        end_path = paths / dir
        ct = bcolz.open(end_path, mode="r")
        daily_data = ct.todataframe()
        daily_data = daily_data.assign(date=daily_data["dt"].apply(lambda x: str(to_datetime(str(x)))[0:10]))
        try:
            del daily_data["factor"]
        except Exception:
            continue
        # is st data
        is_st_data = open_is_st_file(dir)
        # factor data
        factor_data = open_factor_file(dir)
        # merge data
        result = pd.merge(daily_data, is_st_data, on=["symbol", "date"])
        result_data = pd.merge(result, factor_data, on=["symbol", "date"])
        del result_data["date"]
        one_path = write_path / dir
        result_data.sort_values("dt")
        _ctable_kwargs = {"expectedlen": 4096 * 250 * 20}
        ctb = bcolz.ctable.fromdataframe(
            result_data, rootdir=one_path, mode="w", **_ctable_kwargs
        )
        ctb.flush()


def open_factor_file(dir: str):
    global path
    end_path = Path(path) / "factor"
    end_path = end_path / dir
    ct = bcolz.open(end_path, mode="r")
    factor_data = ct.todataframe()
    del factor_data["pre_factor"]
    factor_data = factor_data.rename(columns={"dt": "date", "post_factor": "factor"})
    return factor_data


def open_is_st_file(dir: str):
    global path
    end_path = Path(path) / "st"
    end_path = end_path / dir
    ct = bcolz.open(end_path, mode="r")
    is_st_data = ct.todataframe()
    is_st_data = is_st_data.rename(columns={"dt": "date"})
    return is_st_data


if __name__ == "__main__":
    open_minute_file()




